import cv2
from ultralytics import YOLO


def main():
    # 不写路径就默认从官网上下载 yolov8n.pt
    model = YOLO('models/yolov8n.pt', verbose=True)
    cap = cv2.VideoCapture(0)

    if not cap.isOpened():
        raise RuntimeError('Can not open video device(s)')

    try:
        while True:
            ret, frame = cap.read()
            if ret:
                # Flip the frame horizontally
                frame = cv2.flip(frame, 1)

                model.predict(
                    frame,
                    show=True,  # 显示预测框
                    # save=True,  # 保存预测结果
                    # save_crop=True  # 将预测框中的图像裁剪并保存
                )
            # 按 q 退出
            # & 0xFF 使 waitkey的返回值不超过范围 [0-255]
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    except KeyboardInterrupt as e:
        print(e)
    finally:
        cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
